import os

from langchain.agents import create_structured_chat_agent, AgentExecutor
from langchain.memory import ConversationBufferMemory
from langchain_community.agent_toolkits import create_sql_agent
from langchain_community.utilities import OpenWeatherMapAPIWrapper, SQLDatabase
from langchain_core.tools import tool, Tool
from langchain_openai import ChatOpenAI
from langchain import hub

os.environ["OPENWEATHERMAP_API_KEY"] = "879544e7845335ee2f4d3168fd81cd55"

db = SQLDatabase.from_uri("sqlite:///../my_db.db")

key = 'sk-7wnDma9l5GuVbq38B3C07f50290147148a0809B117A1C1Ad'
model = ChatOpenAI(model="gpt-3.5-turbo",
                   openai_api_key=key,
                   openai_api_base="https://api.aigc369.com/v1",
                   temperature=0)

db_agent_executor = create_sql_agent(
    db=db, llm=model, agent_type="openai-tools", verbose=True, handle_parsing_errors=True
)


@tool
def get_weather(city: str) -> str:
    """天气情况获取工具,当你需要获取天气时，请使用此工具，
    情况一：如果用户问题中没有给出具体的城市消息，那么请先让用户给出具体的城市再使用此工具
    Question：当前天气如何？
    Thought：用户没有给出查询哪个城市的信息
    Action：我需要知道您需要返回你想知道天气的城市名称！
    (Question/Thought/Action 重复调用，直到获取到最后的结果)、


    情况二：如果用户问题中的城市输入格式不符合使用条件，比如
    Question：广州天气条件如何
    Thought：用户现在查询广州的天气情况，但是输入为广州，而不是'GUANGZHOU,CN',所以不满足查询条件
    Action：将广州转变为GUANGZHOU,CN进行查询
    (Question/Thought/Action 重复调用，直到获取到最后的结果)、

    ####
    BEGIN!!!
    Question:
    """
    if "广州" == city:
        return "GUANGZHOU, CN"
    if "北京" == city:
        return "BEIJING, CN"
    if "上海" == city:
        return "SHANGHAI, CN"
    weather = OpenWeatherMapAPIWrapper()
    weather_data = weather.run(city)
    return weather_data


tools = [
    get_weather,
    Tool(
        name="数据库查询工具",
        description="""
            当你本身无法获取到问题追定的数据信息的时候，你可以使用此数据库查询工具，将用户的问题转为SQL后，进行查询操作，最后工具会返回一个数据查询结果
        """,
        func=db_agent_executor.invoke
    )
]

prompt = hub.pull("hwchase17/structured-chat-agent")

agent = create_structured_chat_agent(
    llm=model,
    prompt=prompt,
    tools=tools
)

memory = ConversationBufferMemory(
    return_messages=True,
    memory_key="chat_history"
)

agent_executor = AgentExecutor.from_agent_and_tools(
    agent=agent,
    tools=tools,
    memory=memory,
    verbose=True,
    handle_parsing_errors=True
)

while True:
    input_text = input("请输入你的问题:")
    print(agent_executor.invoke({"input": input_text}))
    # print("=====================================")
# print(agent_executor.invoke({"input": "张三在哪个城市,天气如何?"}))
